#ifndef _WHEELTIMER_H_
#define _WHEELTIMER_H_

#include <thread>
#include <mutex>
#include <atomic>
#include <condition_variable>
#include <chrono>
#include <iostream>
#include <vector>
#include <functional>
#include <set>
#include <unordered_map>

using namespace std::chrono_literals;

class WheelTimer {
public:
    using TickClock = std::chrono::steady_clock;
    using TickDuration = TickClock::duration;
    using TickCount = int32_t;
    using TaskID = int16_t;
    using TaskExpire = int64_t;
    using TaskEntity = std::function<void(void)>;
    using WheelCount = int32_t;

private:
    struct BucketNode {
        TaskID id;
        TaskExpire expire;
        WheelCount turn;
        TaskEntity func;

        BucketNode(TaskID id_, TaskExpire expire_, WheelCount turn_, TaskEntity func_)
        : id(id_), expire(expire_), turn(turn_), func(std::move(func_)) {}

        bool operator<(const WheelTimer::BucketNode& other) const noexcept {
            if (this->expire == other.expire) {
                return this->id < other.id;
            }
            return this->expire < other.expire;
        }

        bool operator==(const WheelTimer::BucketNode& other) const noexcept {
            return this->id == other.id;
        }
    };

    struct WheelBucket {
        std::set<BucketNode> nodes;
    };

public:
    explicit WheelTimer(TickDuration tickDuration_ = 1s, TickCount tickCount_ = 512)
    : tickDuration(tickDuration_), tickCount(tickCount_), stop(false), tickIndex(0),
    counter(0), numTurns(0), buckets(tickCount), worker(&WheelTimer::run, this)
    {

    }

    ~WheelTimer();

    // 添加定时任务
    template<class Rep, class Period, class Function, class ... Args>
    TaskID addTask(const std::chrono::duration<Rep, Period>& d, Function&& f, Args&& ...args) {
        std::lock_guard<std::mutex> lk(wheelMutex);
        TaskID id = counter++;
        TaskExpire expire = std::chrono::nanoseconds(d).count();
        auto ticks = static_cast<TickCount>(expire / tickDuration.count());
        WheelCount turn = numTurns + ticks / tickCount;
        auto task = std::bind(std::forward<Function>(f), std::forward<Args>(args)...);
        std::function<void(void)> func = [task]{ task(); };
        TickCount index = (tickIndex + (ticks % tickCount)) % tickCount;
        auto& st = buckets[index].nodes;
        auto p = st.emplace(id, expire, turn, func);
        taskMap.emplace(id, std::make_pair(index, p.first));
        return id;
    }

    // 取消定时任务
    bool cancelTask(TaskID taskId);

    // 检查定时任务是否还未执行
    bool checkTask(TaskID taskId);

private:
    void run();

private:
    TickDuration tickDuration;
    TickCount tickCount;
    std::atomic<bool> stop;
    TickCount tickIndex;
    TaskID counter;
    WheelCount numTurns;

    std::thread worker;
    std::mutex wheelMutex;
    std::condition_variable wheelCond;
    std::vector<WheelBucket> buckets;
    std::unordered_map<TaskID, std::pair<TickCount, std::set<BucketNode>::iterator>> taskMap;
};


#endif //WHEELTIMER_WHEELTIMER_H
